{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "last_runtime": {
        "build_target": "//third_party/py/jax_triton/google/pallas_tpu:notebook",
        "kind": "private"
      }
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# Pallas Core-specific Programming"
      ],
      "metadata": {
        "id": "YIt0Za36LYg9"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "In this guide, we explore using `pl.core_map` to write Pallas kernels. Compared with `pallas_call`, `core_map` offers a few key characteristics:\n",
        "\n",
        "* **Per-core level programming**: You write code for an TPU/GPU core, not for a JAX device. This gives you full control over what runs on every core, or how cores communicate and distribute work among one another.\n",
        "\n",
        "* **Collectives**: `core_map` explicitly models physical cores, so inter-core communication can be expressed safely.\n",
        "\n",
        "* **Platform generic**: `core_map` programming model works for TPU (TensorCore and SparseCore) and GPU with minimal boilerplate changes."
      ],
      "metadata": {
        "id": "khDWSc7aOVts"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "This guide focuses on TPU. For how to use `core_map` on GPU to achieve higher thread flexibility, check out our [Pallas GPU `core_map` tutorial](https://docs.jax.dev/en/latest/pallas/gpu/reference.html#using-core-map)."
      ],
      "metadata": {
        "id": "i8pl0CLqTVvL"
      }
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Environment setup\n",
        "\n",
        "Modern accelerators often have multiple cores under a device. For recent TPU chips (v4, v5p), every JAX device may contains 2 TensorCores (aka. a [Megacore](https://cloud.google.com/tpu/docs/system-architecture-tpu-vm#chips)). Some TPUs (v5p, v6e, 7x) also contain [SparseCores](https://openxla.org/xla/sparsecore#specifications_at_a_glance), each of which consists of many subcores.\n",
        "\n",
        "This guide was written on a v5p chip, which contains 4 devices (2 TensorCores each) and 4 SparseCores, each with 16 subcores."
      ],
      "metadata": {
        "id": "bsOPXdJkzC-x"
      }
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "id": "14PNaMVsLUur",
        "executionInfo": {
          "status": "ok",
          "timestamp": 1764795546418,
          "user_tz": 480,
          "elapsed": 2087,
          "user": {
            "displayName": "Ivy Zheng",
            "userId": "15297372265856137303"
          }
        },
        "outputId": "01976bb1-2f2f-40e9-ca23-f0e480a82ab3"
      },
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Running on 4 TPU v5p devices.\n"
          ]
        }
      ],
      "source": [
        "from functools import partial\n",
        "\n",
        "import jax\n",
        "from jax.sharding import NamedSharding\n",
        "from jax.experimental import pallas as pl\n",
        "from jax.experimental.pallas import tpu as pltpu\n",
        "from jax.experimental.pallas import tpu_sc as plsc\n",
        "import jax.numpy as jnp\n",
        "import numpy as np\n",
        "\n",
        "\n",
        "num_devices = jax.local_device_count()\n",
        "assert num_devices > 1, \"Please run this notebook with more than one device.\"\n",
        "\n",
        "tpu_info = pltpu.get_tpu_info()  # This notebook only runs on TPU.\n",
        "print(f\"Running on {num_devices} TPU {tpu_info.chip_version} devices.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "In addition to the typical TPU device mesh, you need to make a mesh of cores. Consider this as an addition dimension called `core`, with length 2, in addition to the 4-device mesh you work with. That is 8 cores in total."
      ],
      "metadata": {
        "id": "3f0XEhaYnGyk"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# Mesh of devices\n",
        "mesh = jax.make_mesh((jax.device_count(),), ('device',))\n",
        "print(mesh)\n",
        "\n",
        "# Mesh of cores, within a JAX device\n",
        "tc_mesh = pltpu.create_tensorcore_mesh('core')\n",
        "print(tc_mesh)\n",
        "\n",
        "num_devices = mesh.size\n",
        "num_cores = len(tc_mesh.devices)\n",
        "print(f\"There are {num_devices} devices, and {num_cores} cores each.\")"
      ],
      "metadata": {
        "id": "jr5MARD-mIlC",
        "executionInfo": {
          "status": "ok",
          "timestamp": 1764795546665,
          "user_tz": 480,
          "elapsed": 57,
          "user": {
            "displayName": "Ivy Zheng",
            "userId": "15297372265856137303"
          }
        },
        "outputId": "1ea63c2f-3aec-4cdd-9674-d0e2df32460c"
      },
      "execution_count": 2,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "Mesh('device': 4, axis_types=(Explicit,))\n",
            "TensorCoreMesh(devices=array([TensorCore(id=0), TensorCore(id=1)], dtype=object), axis_names=('core',))\n",
            "There are 4 devices, and 2 cores each.\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "## A simple per-core kernel\n",
        "\n",
        "`pl.core_map` allows you to write per-core local code, just as `jax.shard_map` allows you to write per-device code.\n",
        "\n",
        "In the example kernel below, each core has its own VMEM and semaphore allocations. As with normal kernel, you can initiate copies between HBM and VMEM refs using `pltpu.async_copy`.\n",
        "\n",
        "**Communication between cores**\n",
        "\n",
        "Before communicating between cores, it is good practice to perform a barrier (using `pltpu.semaphore_signal`) to ensure resources have been allocated and both cores are at the same point during the program.\n",
        "\n",
        "Once the cores are synchronized, use `pltpu.make_async_remote_copy` to send data between them. The `device_id` keyword argument generically allows sending to any core on any device, but if you just pass in `{'core': other_core_id}`, it will perform a intra-device inter-core copy (the other axis names are held constant).\n"
      ],
      "metadata": {
        "id": "CYxwiULfndlh"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "# This runs on every core\n",
        "def swap_cores_kernel(in_hbm, out_hbm,\n",
        "                      in_vmem, scratch_vmem, out_vmem,\n",
        "                      sem, send_sem, recv_sem):\n",
        "  core_index = jax.lax.axis_index('core')\n",
        "  num_cores = jax.lax.axis_size('core')\n",
        "  slc_size = in_hbm.shape[-1] // num_cores\n",
        "  slc = pl.ds(core_index * slc_size, slc_size)\n",
        "\n",
        "  # Copy in a core-dependent slice of the input\n",
        "  pltpu.async_copy(in_hbm.at[:, slc], in_vmem, sem).wait()\n",
        "\n",
        "  # A barrier to make sure all cores have entered run_scoped.\n",
        "  # You won't need this if not doing inter-core communications.\n",
        "  dst_core = (core_index + 1) % num_cores\n",
        "  sem0 = pltpu.get_barrier_semaphore()\n",
        "  pltpu.semaphore_signal(sem0, 1, device_id={'core': dst_core})\n",
        "  pltpu.semaphore_wait(sem0, 1)\n",
        "\n",
        "  # Swap data between core 0 and core 1\n",
        "  the_copy = pltpu.make_async_remote_copy(\n",
        "      in_vmem, scratch_vmem, send_sem, recv_sem, device_id={'core': dst_core},\n",
        "  )\n",
        "  the_copy.start()\n",
        "  the_copy.wait()\n",
        "\n",
        "  # Core-local compute\n",
        "  out_vmem[...] = scratch_vmem[...] * 2\n",
        "\n",
        "  # Copy out the output\n",
        "  pltpu.async_copy(out_vmem, out_hbm.at[:, slc], sem).wait()\n"
      ],
      "metadata": {
        "id": "GkGRT2HRJOUU",
        "executionInfo": {
          "status": "ok",
          "timestamp": 1764795546946,
          "user_tz": 480,
          "elapsed": 53,
          "user": {
            "displayName": "Ivy Zheng",
            "userId": "15297372265856137303"
          }
        }
      },
      "execution_count": 3,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "Once you have the local kernel:\n",
        "\n",
        " * Start your top-level JAX code with HBM refs, and allocate output refs if needed.\n",
        "\n",
        " * Use `pl.core_map`, which takes the TensorCore mesh, to start per-core programming.\n",
        "\n",
        "    * You will need `collective_id` for the barrier semaphore.\n",
        "\n",
        " * Inside `pl.core_map`, invoke `pl.run_scoped` to allocate per-core scratch spaces (VMEM and semaphores) and run the local kernel."
      ],
      "metadata": {
        "id": "2T0tSkFmoFLI"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "input_shape = (32, 256)\n",
        "local_vmem_shape = (32 // num_devices, 256 // num_cores)\n",
        "in_spec = jax.P('device', None)\n",
        "sharding = NamedSharding(mesh, in_spec)\n",
        "\n",
        "@jax.jit\n",
        "@partial(jax.shard_map, mesh=mesh, in_specs=in_spec, out_specs=in_spec,\n",
        "         check_vma=False)\n",
        "def swap_cores(x):\n",
        "  # Get buffers out of the input and output\n",
        "  x_hbm_ref = jax.new_ref(x)\n",
        "  o_hbm_ref = jax.new_ref(jax.lax.empty(x.shape, x.dtype))\n",
        "\n",
        "  @pl.core_map(tc_mesh, compiler_params=pltpu.CompilerParams(collective_id=0))\n",
        "  def _():\n",
        "    pl.run_scoped(\n",
        "        partial(swap_cores_kernel, x_hbm_ref, o_hbm_ref),\n",
        "        *([pltpu.VMEM(local_vmem_shape, x.dtype)] * 3),  # VMEM allocations\n",
        "        *([pltpu.SemaphoreType.DMA] * 3),                # semaphores\n",
        "    )\n",
        "  return o_hbm_ref[...]\n",
        "\n",
        "\n",
        "x = jax.random.normal(jax.random.key(0), input_shape, jnp.float32)\n",
        "x = jax.device_put(x, sharding)\n",
        "y = swap_cores(x)\n",
        "\n",
        "np.testing.assert_array_equal(y[:, 128:], x[:, :128] * 2)\n",
        "np.testing.assert_array_equal(y[:, :128], x[:, 128:] * 2)"
      ],
      "metadata": {
        "id": "KT6zkEKi1Sbc",
        "executionInfo": {
          "status": "ok",
          "timestamp": 1764795548996,
          "user_tz": 480,
          "elapsed": 1800,
          "user": {
            "displayName": "Ivy Zheng",
            "userId": "15297372265856137303"
          }
        }
      },
      "execution_count": 4,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "### Save the boilerplate\n",
        "\n",
        "You can use the `pl.kernel` decorator to wrap boilerplate such as `core_map`, `run_scoped`, and output buffer allocation.\n",
        "\n",
        "Note that this should run inside any `jax.shard_map` you may have at the top level."
      ],
      "metadata": {
        "id": "dLV8sKa4HuSX"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "@jax.jit\n",
        "@partial(jax.shard_map, mesh=mesh, in_specs=in_spec, out_specs=in_spec, check_vma=False)\n",
        "def swap_cores(x):\n",
        "  scratch_shapes = [pltpu.VMEM(local_vmem_shape, x.dtype)] * 3 + [pltpu.SemaphoreType.DMA] * 3\n",
        "  return pl.kernel(swap_cores_kernel, out_shape=x, mesh=tc_mesh,\n",
        "                   scratch_shapes=scratch_shapes,\n",
        "                   compiler_params=pltpu.CompilerParams(collective_id=0))(x)\n",
        "\n",
        "y = swap_cores(x)\n",
        "np.testing.assert_array_equal(y[:, 128:], x[:, :128] * 2)\n",
        "np.testing.assert_array_equal(y[:, :128], x[:, 128:] * 2)"
      ],
      "metadata": {
        "id": "7cHnsRHPHyfH",
        "executionInfo": {
          "status": "ok",
          "timestamp": 1764795549347,
          "user_tz": 480,
          "elapsed": 106,
          "user": {
            "displayName": "Ivy Zheng",
            "userId": "15297372265856137303"
          }
        }
      },
      "execution_count": 5,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Pipelining with `core_map`\n",
        "\n",
        "Note that the kernel above only does simple copies and compute, without automatic pipelining via Pallas `grid` and `BlockSpec`. To do pipelining inside `core_map`, use `pltpu.emit_pipeline` inside the core-local kernel.\n",
        "\n",
        "**Automatically parallelize work amongst cores**\n",
        "\n",
        "The simple way is to annotate a block axis as `pltpu.PARALLEL`, and Pallas will automatically parallelize work along this axis. Both `pl.pallas_call` and `pltpu.emit_pipeline` supports this, via arguments `core_axis` and `dimension_semantics`. The `pallas_call` example is [in another guide](https://docs.jax.dev/en/latest/pallas/tpu/pipelining.html#tpus-in-megacore-configuration), and the `emit_pipeline` case is shown below.\n",
        "\n",
        "When the `PARALLEL` annotation is provided, the corresponding grid dimension will be logically split and executed on separate cores. (The exact semantics of which grid dimensions are executed on which core is guaranteed).\n",
        "\n",
        "**Scratch shapes allocation**\n",
        "\n",
        "Note that in the example below, the top level `pl.run_scoped` (wrapped inside `kernel`) did not allocate any VMEM scratch buffers. Instead, `pltpu.emit_pipeline` allocates its own scratch buffers in VMEM and use them for its multiple buffering.\n"
      ],
      "metadata": {
        "id": "4-G--Wnysdjs"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "def add_one_body(in_vmem, out_vmem):\n",
        "  out_vmem[...] = in_vmem[...] + 1\n",
        "\n",
        "input_shape = (1024, 1024)\n",
        "in_spec = jax.P('device', None)\n",
        "\n",
        "def add_one_kernel(x_hbm_ref, o_hbm_ref):\n",
        "  in_shape = x_hbm_ref.shape\n",
        "  pltpu.emit_pipeline(\n",
        "      add_one_body,\n",
        "      grid=(in_shape[0] // 8, in_shape[1] // 128),\n",
        "      in_specs=[pl.BlockSpec(\n",
        "          block_shape=(8, 128), index_map=lambda i, j: (i, j),\n",
        "      )],\n",
        "      out_specs=[pl.BlockSpec(\n",
        "          block_shape=(8, 128), index_map=lambda i, j: (i, j),\n",
        "      )],\n",
        "      core_axis_name='core',\n",
        "      dimension_semantics=(pltpu.PARALLEL, pltpu.ARBITRARY),\n",
        "  )(x_hbm_ref, o_hbm_ref)\n",
        "\n",
        "\n",
        "@jax.jit\n",
        "@partial(jax.shard_map, mesh=mesh, in_specs=in_spec, out_specs=in_spec, check_vma=False)\n",
        "def add_one(x):\n",
        "  return pl.kernel(add_one_kernel, out_shape=x, mesh=tc_mesh, scratch_shapes=[])(x)\n",
        "\n",
        "\n",
        "x = jax.random.normal(jax.random.key(0), input_shape, jnp.float32)\n",
        "x = jax.device_put(x, NamedSharding(mesh, in_spec))\n",
        "y = add_one(x)\n",
        "\n",
        "np.testing.assert_array_equal(y, x + 1)"
      ],
      "metadata": {
        "id": "xUMRPLxb1rEH",
        "executionInfo": {
          "status": "ok",
          "timestamp": 1764795550106,
          "user_tz": 480,
          "elapsed": 518,
          "user": {
            "displayName": "Ivy Zheng",
            "userId": "15297372265856137303"
          }
        }
      },
      "execution_count": 6,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Scalar prefetch\n",
        "\n",
        "The code below extends the kernel above but uses [scalar prefetch and dynamic block indexing](https://docs.jax.dev/en/latest/pallas/tpu/sparse.html) to select a specific sub-slice of the input.\n",
        "\n",
        "This involves pre-allocating an SMEM buffer (via the `pl.run_scoped` call inside `kernel`) and populating the buffer using a `sync_copy` before the pipeline starts. Close over the dynamic index value inside the `index_map` to use it.\n",
        "\n",
        "**Manually delegate work amongst cores**\n",
        "\n",
        "The code example below also shows how `core_map` allows you to customize exactly how the work is split between cores, without relying on the automatic API shown above.\n",
        "\n",
        "To achieve that, customize your `index_map` to use the core index to work on different slices on different cores.\n"
      ],
      "metadata": {
        "id": "Cq5rYyvL2Tte"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "input_shape = (1024, 1024)\n",
        "in_spec = jax.P('device', None)\n",
        "output_shape = (1024, 512)\n",
        "\n",
        "def indexed_add_one_kernel(in_refs, out_refs, i_smem_ref):\n",
        "  (x_hbm_ref, i_hbm_ref), o_hbm_ref = in_refs, out_refs\n",
        "  in_shape = x_hbm_ref.shape\n",
        "  pltpu.sync_copy(i_hbm_ref, i_smem_ref)\n",
        "\n",
        "  core_idx = jax.lax.axis_index('core')\n",
        "  core_slc_size = in_shape[0] // num_cores\n",
        "  i_map = lambda i: core_idx * core_slc_size // 8 + i  # split work among cores\n",
        "  j_map = lambda j: i_smem_ref[0] // 128 + j           # use the prefetched offset\n",
        "\n",
        "  pltpu.emit_pipeline(\n",
        "      add_one_body,\n",
        "      grid=(core_slc_size // 8, output_shape[1] // 128),\n",
        "      in_specs=[pl.BlockSpec(\n",
        "          block_shape=(8, 128), index_map=lambda i, j: (i_map(i), j_map(j)),\n",
        "      )],\n",
        "      out_specs=[pl.BlockSpec(\n",
        "          block_shape=(8, 128), index_map=lambda i, j: (i_map(i), j),\n",
        "      )]\n",
        "  )(x_hbm_ref, o_hbm_ref)\n",
        "\n",
        "\n",
        "@jax.jit\n",
        "@partial(jax.shard_map, mesh=mesh,\n",
        "         in_specs=(in_spec, jax.P()), out_specs=in_spec, check_vma=False)\n",
        "def indexed_add_one(x, index):\n",
        "  out_shape = jax.ShapeDtypeStruct((x.shape[0], x.shape[1] // 2), x.dtype)\n",
        "  return pl.kernel(indexed_add_one_kernel,\n",
        "                   out_shape=out_shape, mesh=tc_mesh,\n",
        "                   scratch_shapes=[pltpu.SMEM((1,), jnp.int32)])((x, index))\n",
        "\n",
        "\n",
        "xs = jax.random.normal(jax.random.key(0), input_shape, jnp.float32)\n",
        "xs = jax.device_put(xs, NamedSharding(mesh, in_spec))\n",
        "idx = 256\n",
        "y = indexed_add_one(xs, jnp.array([idx]))\n",
        "\n",
        "np.testing.assert_array_equal(y, xs[:, idx:(idx+512)] + 1)"
      ],
      "metadata": {
        "id": "SE8pTStHeSWB",
        "executionInfo": {
          "status": "ok",
          "timestamp": 1764795550778,
          "user_tz": 480,
          "elapsed": 378,
          "user": {
            "displayName": "Ivy Zheng",
            "userId": "15297372265856137303"
          }
        }
      },
      "execution_count": 7,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Mapping over SparseCores\n",
        "\n",
        "TPU v5p contains 4 [SparseCores](https://openxla.org/xla/sparsecore), which are specialized for sparse memory access and operations. This guide will not dive into the full capabilities of SparseCore, but rather show how to run a program on SparseCore with the same semantics and minimal changes from the TensorCore code.\n",
        "\n",
        "Start with knowing the basic SparseCore specs of your chip, and create a `VectorSubcoreMesh` for vector operations. Note that each SparseCore has 16 (or other number) subcores on TPU v5p, and `core_map` will run your code SPMD on each of them."
      ],
      "metadata": {
        "id": "B8qeo-4A2KRm"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "sc_info = pltpu.get_tpu_info().sparse_core\n",
        "assert sc_info is not None\n",
        "print(sc_info)\n",
        "\n",
        "sc_mesh = plsc.VectorSubcoreMesh(\n",
        "    core_axis_name=\"core\", subcore_axis_name=\"subcore\",\n",
        "    num_cores=sc_info.num_cores\n",
        ")\n",
        "sc_num_cores = sc_info.num_cores\n",
        "sc_num_subcores = sc_info.num_subcores"
      ],
      "metadata": {
        "id": "AHurx-yyYVvs",
        "executionInfo": {
          "status": "ok",
          "timestamp": 1764795551102,
          "user_tz": 480,
          "elapsed": 55,
          "user": {
            "displayName": "Ivy Zheng",
            "userId": "15297372265856137303"
          }
        },
        "outputId": "aa4a45da-dd9a-4f57-de1a-bc9b5872b2df"
      },
      "execution_count": 8,
      "outputs": [
        {
          "output_type": "stream",
          "name": "stdout",
          "text": [
            "SparseCoreInfo(num_cores=4, num_subcores=16, num_lanes=8)\n"
          ]
        }
      ]
    },
    {
      "cell_type": "markdown",
      "source": [
        "The code below is very similar to the `add_one_kernel` we wrote earlier, except for a few differences:\n",
        "\n",
        "1. You need to split the work amongst all subcores, so a few lines to compute the specific slice for each subcore.\n",
        "\n",
        "1. SparseCore register computation allows smaller slices (`4x16` max for int32), so you need nested loops to iterate the slice during computation phase."
      ],
      "metadata": {
        "id": "n2_dfsUWFgwU"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "input_shape = (4096, 128)\n",
        "SC_REG_OP_SHAPE = (4, 16)\n",
        "\n",
        "def sc_add_one_body(in_vmem, out_vmem):\n",
        "  @pl.loop(0, in_vmem.shape[0], step=SC_REG_OP_SHAPE[0])\n",
        "  def _reg_loop_0(c0):\n",
        "    @pl.loop(0, in_vmem.shape[1], step=SC_REG_OP_SHAPE[1])\n",
        "    def _reg_loop_1(c1):\n",
        "      slc = (pl.ds(c0, SC_REG_OP_SHAPE[0]), pl.ds(c1, SC_REG_OP_SHAPE[1]))\n",
        "      out_vmem[slc] = in_vmem[slc] + 1\n",
        "\n",
        "\n",
        "def sc_add_one_kernel(x_hbm_ref, o_hbm_ref):\n",
        "  in_shape = x_hbm_ref.shape\n",
        "  core_idx = jax.lax.axis_index('core')\n",
        "  subcore_idx = jax.lax.axis_index(\"subcore\")\n",
        "  cm_idx = core_idx * sc_num_subcores + subcore_idx  # index on the core_map\n",
        "  slc_size = in_shape[0] // (sc_num_subcores * sc_num_cores)\n",
        "  index_map = lambda i, j: (\n",
        "      pl.ds(pl.multiple_of(cm_idx * slc_size + i * 8, 8), 8), j)\n",
        "\n",
        "  pltpu.emit_pipeline(\n",
        "      sc_add_one_body,\n",
        "      grid=(slc_size // 8, in_shape[1] // 128),\n",
        "      in_specs=[pl.BlockSpec(\n",
        "          block_shape=(pl.BoundedSlice(8), 128), index_map=index_map,\n",
        "      )],\n",
        "      out_specs=[pl.BlockSpec(\n",
        "          block_shape=(pl.BoundedSlice(8), 128), index_map=index_map,\n",
        "      )]\n",
        "  )(x_hbm_ref, o_hbm_ref)\n",
        "\n",
        "\n",
        "@jax.jit\n",
        "@partial(jax.shard_map, mesh=mesh, in_specs=in_spec, out_specs=in_spec, check_vma=False)\n",
        "def sc_add_one(x):\n",
        "  return pl.kernel(sc_add_one_kernel, out_shape=x, mesh=sc_mesh, scratch_shapes=[])(x)\n",
        "\n",
        "\n",
        "x = jax.random.randint(jax.random.key(0), input_shape, 0, 64, jnp.int32)\n",
        "x = jax.device_put(x, NamedSharding(mesh, in_spec))\n",
        "y = sc_add_one(x)\n",
        "\n",
        "np.testing.assert_array_equal(y, x + 1)"
      ],
      "metadata": {
        "id": "6fNShx6k2kxi",
        "executionInfo": {
          "status": "ok",
          "timestamp": 1764795552411,
          "user_tz": 480,
          "elapsed": 1117,
          "user": {
            "displayName": "Ivy Zheng",
            "userId": "15297372265856137303"
          }
        }
      },
      "execution_count": 9,
      "outputs": []
    }
  ]
}